from torch import nn


def conv_unit(in_channels, out_channels, kernel_size, stride, use_bn=False, act_type=None):
    pad = kernel_size // 2
    layers = []
    layers.append(nn.Conv2d(in_channels, out_channels, kernel_size, stride, pad, bias=False))
    if use_bn:
        layers.append(nn.BatchNorm2d(out_channels))
    if act_type == 'prelu':
        layers.append(nn.PReLU())
    elif act_type == 'relu':
        layers.append(nn.ReLU())
    return layers
